import numpy as np
import re
import random


def text_parse(input_string: str) -> list:
    """
    文章分词
    :param input_string: 输入的文章
    :return: 分词列表
    """
    list_of_tokens = re.split(r'\W+', input_string)
    return [token.lower() for token in list_of_tokens if len(list_of_tokens) > 2]


def create_vocab_list(doc_list):
    """
    创建词汇表
    :param doc_list: 文件单词列表
    :return: 词汇表
    """
    vocab_set = set([])
    for document in doc_list:
        vocab_set = vocab_set | set(document)
    return list(vocab_set)


def set_of_word2_vec(vocab_list, input_set):
    """
    词汇表输入文本向量化
    :param vocab_list: 词汇表
    :param input_set: 输入文章
    :return: 向量化的词汇表
    """
    return_vec = [0] * len(vocab_list)
    for word in input_set:
        if word in vocab_list:
            return_vec[vocab_list.index(word)] = 1
    return return_vec


def train_nb(train_mat, train_class):
    """
    计算训练集的 P0(D|h) 和 P1(D|h) 以及 P1(h)
    :param train_mat: 训练集向量化词汇表
    :param train_class: 训练集类别
    :return: P0(D|h)、P1(D|h)、P1(h)
    """
    num_examples = len(train_mat)
    num_words = len(train_mat[0])
    p1_class = sum(train_class) / float(num_examples)
    p0_num = np.ones(num_words)     # 不能乘零
    p1_num = np.ones(num_words)
    p0_den = 2
    p1_den = 2

    for i in range(num_examples):
        if train_class[i] == 1:
            p1_num += train_mat[i]
            p1_den += sum(train_mat[i])
        else:
            p0_num += train_mat[i]
            p0_den += sum(train_mat[i])

    p1_vec = np.log(p1_num / p1_den)
    p0_vec = np.log(p0_num / p0_den)
    return p0_vec, p1_vec, p1_class


def classify_nb(word_vec, p0_vec, p1_vec, p1_class):
    """
    进行 log 相加计算邮件概率
    :param word_vec: 向量化词汇表
    :param p0_vec: P0(D|h)
    :param p1_vec: P1(D|h)
    :param p1_class: P1(h)
    :return: 预测类别 0 | 1
    """
    # word_vec * p1_vec 去除垃圾邮箱中未出现词汇的频率
    p1 = np.log(p1_class) + sum(word_vec * p1_vec)
    p0 = np.log(1 - p1_class) + sum(word_vec * p0_vec)
    return 1 if p1 > p0 else 0


def spam():
    """
    主函数
    """
    # 读取文件，并对文件进行 分词 和 分类 。
    doc_list = []
    class_list = []
    for i in range(1, 26):
        wordlist = text_parse(open(f'email/spam/{i}.txt', 'r').read())
        doc_list.append(wordlist)
        class_list.append(1)    # 1 垃圾邮件

        wordlist = text_parse(open(f'email/ham/{i}.txt', 'r').read())
        doc_list.append(wordlist)
        class_list.append(0)    # 0 正常邮件

    # 创建词汇表
    vocab_list = create_vocab_list(doc_list)

    # 随机切分数据集索引为 训练集 和 测试集
    train_set = list(range(50))
    test_set = []
    for i in range(10):
        # random.uniform(a, b) 根据四舍五入，获取 [a, b) 或 [a, b] 范围内的随机数。
        rand_index = int(random.uniform(0, len(train_set)))
        test_set.append(train_set[rand_index])
        del train_set[rand_index]

    # 训练集向量化并获取标签值
    train_mat = []
    train_class = []
    for doc_index in train_set:
        train_mat.append(set_of_word2_vec(vocab_list, doc_list[doc_index]))
        train_class.append(class_list[doc_index])

    # 统计计算训练出的 正常邮件词汇表频率P0(D|h) 和 垃圾邮件词汇表频率P1(D|h) 以及 垃圾邮件概率P1(h)
    p0_vec, p1_vec, p1_class = train_nb(np.array(train_mat), np.array(train_class))

    # 预测测试集
    error_count = 0
    for doc_index in test_set:
        word_vec = set_of_word2_vec(vocab_list, doc_list[doc_index])
        predict = classify_nb(np.array(word_vec), p0_vec, p1_vec, p1_class)
        if predict != class_list[doc_index]:
            error_count += 1
    print('当前10个测试样本，错误 :', error_count, '个')


if __name__ == '__main__':
    spam()
